import os
import numpy as np
from sklearn.model_selection import train_test_split

PROCESSED_DATA_PATH = '../data/processed/'  # 预处理后音频特征文件的路径
TRAIN_DATA_SAVE_PATH = '../data/train_data.npy'
VAL_DATA_SAVE_PATH = '../data/val_data.npy'
TEST_DATA_SAVE_PATH = '../data/test_data.npy'
TRAIN_LABELS_SAVE_PATH = '../data/train_labels.npy'
VAL_LABELS_SAVE_PATH = '../data/val_labels.npy'
TEST_LABELS_SAVE_PATH = '../data/test_labels.npy'

TRAIN_SIZE = 0.7  # 训练集所占比例
VAL_SIZE = 0.2  # 验证集所占比例（相对于原始数据集）
TEST_SIZE = 0.1  # 测试集所占比例（相对于原始数据集）
RANDOM_STATE = 42  # 随机种子，保证结果可重复

LABELS = {'非病理性': 0, '哮喘': 1, '咽炎': 2, '支气管炎': 3, '百日咳': 4}  # 标签映射


def load_features_and_labels():
    """
    加载预处理后的特征文件和标签。
    Returns:
        features (numpy array): 特征数据。
        labels (numpy array): 标签数据。
    """
    features = []
    labels = []

    for file_name in os.listdir(PROCESSED_DATA_PATH):
        if file_name.endswith('.npy'):
            feature = np.load(os.path.join(PROCESSED_DATA_PATH, file_name))
            features.append(feature)

            # 使用预定义的标签映射
            label_name = file_name.split('_')[0]
            if label_name in LABELS:
                labels.append(LABELS[label_name])
            else:
                raise ValueError(f"Unknown label: {label_name}")

    features = np.array(features)
    labels = np.array(labels)

    return features, labels


def split_and_save_data(features, labels):
    """
    划分数据集并保存到文件。
    Args:
        features (numpy array): 特征数据。
        labels (numpy array): 标签数据。
    """
    # 首先将数据划分为训练集和临时集（包括验证集和测试集）
    X_train, X_temp, y_train, y_temp = train_test_split(features, labels, train_size=TRAIN_SIZE,
                                                        random_state=RANDOM_STATE)

    # 然后从临时集中再划分出验证集和测试集
    val_size_relative = VAL_SIZE / (VAL_SIZE + TEST_SIZE)
    X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=(1 - val_size_relative),
                                                    random_state=RANDOM_STATE)

    # 保存数据集
    np.save(TRAIN_DATA_SAVE_PATH, X_train)
    np.save(VAL_DATA_SAVE_PATH, X_val)
    np.save(TEST_DATA_SAVE_PATH, X_test)
    np.save(TRAIN_LABELS_SAVE_PATH, y_train)
    np.save(VAL_LABELS_SAVE_PATH, y_val)
    np.save(TEST_LABELS_SAVE_PATH, y_test)


def main():
    features, labels = load_features_and_labels()

    # 调整特征形状，以适应模型的输入格式
    # features.shape为(1432, 141, 173)，调整为（1432, 141, 173, 1）
    features = features[..., np.newaxis]

    split_and_save_data(features, labels)


if __name__ == '__main__':
    main()
